# %%
import torch
from torchvision import datasets, transforms
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

# %%
print("PyTorch Version: ", torch.__version__)
# %%
# 检测cuda是否可用
use_cuda = torch.cuda.is_available()
print("use_cuda:", use_cuda)
# %%
# 设置device变量
if use_cuda:
    device = torch.device("cuda")
else:
    device = torch.device("cpu")
# %%
# 设置对数据处理的逻辑
transform = transforms.Compose([
    # 将图像转化为Tensor张量
    transforms.ToTensor(),
    # 将图像的像素值归一化到[0,1]之间，0.1307和0.3081是MNIST数据集的均值和方差
    transforms.Normalize((0.1307,), (0.3081,))
])
# %%
# 加载MNIST数据集
train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)
# %% md
# ### 计算训练集和测试集的均值和方差
# %%
# 设置数据加载器DataLoader， 顺便设置批次大小和是否打乱数据
# train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=60000, shuffle=True)
# test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=100, shuffle=False)
#
# for batch_idx, data in enumerate(train_loader, 0):
#     inputs, targets = data
#     # view函数会将训练集（60000,1，28，28）转化为（60000，784）
#     x = inputs.view(-1, 28 * 28)
#     x_std = x.std().item()
#     x_mean = x.mean().item()
#
# print("x_std:" + str(x_std))
# print("x_mean:" + str(x_mean))
# %%
# 设置数据加载器DataLoader， 顺便设置批次大小和是否打乱数据
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=128, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=1000, shuffle=False)


# %% md
# ### 通过自定义类来构建模型
# %%
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(784, 128)
        self.dropout = nn.Dropout(0.2)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dropout(x)
        x = self.fc2(x)
        output = F.log_softmax(x, dim=1)
        return output


# %% md
# ### 实例化模型，并将模型发送到device
# %%
model = Net().to(device)


# %% md
# 
# %% md
# ### 定义训练模型的逻辑
# %%
def train_step(data, target, model, optimizer):
    optimizer.zero_grad()
    output = model(data)
    # nll_loss函数会计算输入和目标的负对数似然, nll是 negative log likelihood 的缩写--负对数似然
    # 这里的target是one-hot编码的，所以需要使用nll_loss函数
    loss = F.nll_loss(output, target)
    # 反向传播的本质是计算梯度，并更新模型参数
    loss.backward()
    # 本质是应用梯度更新模型参数
    optimizer.step()
    return loss.item()


# %% md
# ### 定义测试模型的逻辑
# %%
def test_step(data, target, model, test_loss, correct):
    output = model(data)
    # 计算累积的批次损失
    test_loss += F.nll_loss(output, target, reduction='sum').item()
    # 获得对数概率最大的下标， 这里其实是类别号
    pred = output.argmax(dim=1, keepdim=True)  # dim=1表示按列取最大值，keepdim=True表示保持维度
    correct += pred.eq(target.view_as(pred)).sum().item()
    return test_loss, correct


# %% md
# ### 设置训练调参使用的优化器
# %%
optimizer = optim.Adam(model.parameters(), lr=0.001)
# %% md
# ### 训练模型
# %%
EPOCHS = 5

for epoch in range(EPOCHS):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        loss = train_step(data, target, model, optimizer)
        # 每隔10个批次打印信息
        if batch_idx % 10 == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(epoch, batch_idx * len(data),
                                                                           len(train_loader.dataset),
                                                                           100. * batch_idx / len(train_loader),
                                                                           loss))

    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            test_loss, correct = test_step(data, target, model, test_loss, correct)
    test_loss /= len(test_loader.dataset)
    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))
